#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <utility>
#include <vector>
#include <queue>
using namespace std;

using idx_t = int32_t;

struct IndexNSW {
    using Entry = std::pair<float, idx_t>;
    using CandidateQueue = std::priority_queue<Entry, std::vector<Entry>, greater<>>;
    using ResultQueue = std::priority_queue<Entry, std::vector<Entry>, less<>>;

    size_t d_;
    size_t ntotal_;
    size_t m_; /// neighbors to add link

    idx_t entry_point_ = -1;
    std::unique_ptr<float[]> embedding_; /// embeddings of vector
    std::vector<std::vector<idx_t>> neighbors_; /// the neighbors of the vector
    

    explicit IndexNSW(size_t d, 
                      size_t n, 
                      std::unique_ptr<float[]> embedding,
                      size_t m = 16) 
                : d_(d), ntotal_(n), embedding_(std::move(embedding)), m_(m) {}

    void train();


    std::vector<idx_t> search(const float *xq, size_t k);

    const float* get_ith_vector(size_t i) {
        return embedding_.get() + i * d_;
    }

    std::vector<idx_t>& get_neighbors(idx_t i) {
        return neighbors_[i];
    }

};

float fvec_l2(const float *x, const float *y, size_t d);


